import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import OrderedDict


class Bottleneck(nn.Module):
    """resnet bottleneck"""
    expansion = 4

    def __init__(self, in_channel, mid_channel, stride=1, downsample=None):
        if downsample == None:
            assert mid_channel * self.expansion == in_channel, f"in_channel(which is {in_channel})*self.expansion(which is {self.expansion}) != mid_channel(which is {mid_channel})"

        super(Bottleneck, self).__init__()

        self.conv1 = nn.Conv2d(in_channels=in_channel,
                               out_channels=mid_channel,
                               kernel_size=1,
                               stride=1,
                               bias=False)  # squeeze channels
        self.bn1 = nn.BatchNorm2d(mid_channel)

        self.conv2 = nn.Conv2d(in_channels=mid_channel,
                               out_channels=mid_channel,
                               kernel_size=3,
                               stride=stride,
                               bias=False,
                               padding=1)
        self.bn2 = nn.BatchNorm2d(mid_channel)

        self.conv3 = nn.Conv2d(in_channels=mid_channel,
                               out_channels=mid_channel * self.expansion,
                               kernel_size=1,
                               stride=1,
                               bias=False)  # unsqueeze channels
        self.bn3 = nn.BatchNorm2d(mid_channel * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample

    def forward(self, x):
        identity = x
        if self.downsample is not None:
            identity = self.downsample(x)
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        out += identity
        out = self.relu(out)

        return out


class BasicBlock(nn.Module):
    """resnet basicblock"""
    expansion = 1
    def __init__(self, in_channel, out_channel, stride=1, downsample=None):
        if downsample is None:
            assert in_channel * self.expansion == out_channel, " in_channel*self.expansion != out_channel "
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=in_channel,
                               out_channels=out_channel,
                               kernel_size=3,
                               stride=stride,
                               padding=1,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(out_channel)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(in_channels=out_channel,
                               out_channels=out_channel,
                               kernel_size=3,
                               stride=1,
                               padding=1,
                               bias=False)
        self.bn2 = nn.BatchNorm2d(out_channel)
        self.downsample = downsample

    def forward(self, x):
        identity = x
        if self.downsample is not None:
            identity = self.downsample(x)
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        out += identity
        out = self.relu(out)

        return out
